#!/usr/bin/env python3
"""
批量抽帧脚本
--------------------------------
功能:
  - 从指定文件夹批量提取视频帧
  - 支持间隔抽帧、分辨率缩放、多线程处理
用法示例:
  python extract_frames_batch.py -i downloads -o frames --fps 1
  python extract_frames_batch.py -i ./videos --every-n 10 --resize 1280x720 --threads 4
"""

import cv2
import os
import argparse
import concurrent.futures
from tqdm import tqdm

def extract_frames(video_path, output_dir, fps=None, every_n=None, resize=None):
    """从单个视频中抽帧"""
    try:
        video_name = os.path.splitext(os.path.basename(video_path))[0]
        save_dir = os.path.join(output_dir, video_name)
        os.makedirs(save_dir, exist_ok=True)

        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f"[警告] 无法打开视频：{video_path}")
            return 0

        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        video_fps = cap.get(cv2.CAP_PROP_FPS) or 25

        # 计算抽帧间隔
        frame_interval = 1
        if fps:  # 每秒多少帧
            frame_interval = int(round(video_fps / fps))
        elif every_n:
            frame_interval = every_n

        resize_wh = None
        if resize:
            try:
                w, h = map(int, resize.lower().split('x'))
                resize_wh = (w, h)
            except:
                print(f"[警告] 无法解析 resize 参数: {resize}")

        frame_idx, saved = 0, 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            if frame_idx % frame_interval == 0:
                if resize_wh:
                    frame = cv2.resize(frame, resize_wh)
                frame_path = os.path.join(save_dir, f"{video_name}_frame_{frame_idx:06d}.jpg")
                cv2.imwrite(frame_path, frame)
                saved += 1
            frame_idx += 1

        cap.release()
        return saved

    except Exception as e:
        print(f"[错误] 处理视频 {video_path} 时出错：{e}")
        return 0

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--input', required=True, help='输入视频文件夹')
    parser.add_argument('-o', '--output', default='frames', help='输出图像文件夹')
    parser.add_argument('--fps', type=float, default=None, help='每秒抽取多少帧（例如 1 表示每秒 1 帧）')
    parser.add_argument('--every-n', type=int, default=None, help='每隔 N 帧抽取一张（fps 与此二选一）')
    parser.add_argument('--resize', type=str, default=None, help='调整图像大小，如 1280x720')
    parser.add_argument('--threads', type=int, default=1, help='并行线程数')
    args = parser.parse_args()

    os.makedirs(args.output, exist_ok=True)

    # 支持的视频格式
    video_exts = ('.mp4', '.mkv', '.mov', '.avi', '.flv')
    videos = [os.path.join(args.input, f) for f in os.listdir(args.input)
              if f.lower().endswith(video_exts)]

    if not videos:
        print("[错误] 输入文件夹中没有找到视频文件。")
        return

    print(f"共检测到 {len(videos)} 个视频，将输出至：{args.output}")
    print(f"抽帧方式: {'每秒 ' + str(args.fps) + ' 帧' if args.fps else '每 ' + str(args.every_n) + ' 帧'}")
    if args.resize:
        print(f"输出分辨率: {args.resize}")

    results = []
    if args.threads > 1:
        with concurrent.futures.ThreadPoolExecutor(max_workers=args.threads) as executor:
            futures = {executor.submit(extract_frames, v, args.output, args.fps, args.every_n, args.resize): v for v in videos}
            for f in tqdm(concurrent.futures.as_completed(futures), total=len(videos), desc="处理中"):
                results.append((futures[f], f.result()))
    else:
        for v in tqdm(videos, desc="处理中"):
            count = extract_frames(v, args.output, args.fps, args.every_n, args.resize)
            results.append((v, count))

    total_saved = sum(c for _, c in results)
    print(f"\n✅ 抽帧完成，共保存 {total_saved} 张图像。")

if __name__ == '__main__':
    main()
